//
//  ViewController.swift
//  CoreMLDemo
//
//  Created by stoicer on 2024/6/6.
//

import UIKit
import CoreML

class ViewController: UIViewController,UINavigationControllerDelegate {

    @IBOutlet weak var nameLabel: UILabel!
    @IBOutlet weak var probLabel: UILabel!
    @IBOutlet weak var imageView: UIImageView!
    
    fileprivate let imageController = UIImagePickerController()
    let model = ImageRecognizer()
    override func viewDidLoad() {
        super.viewDidLoad()
       
    }


    @IBAction func choiceClick(_ sender: Any) {
        imageController.delegate = self
        self.present(imageController, animated: true)
    }
    
    func buffer(from image: UIImage) -> CVPixelBuffer? {
      let attrs = [kCVPixelBufferCGImageCompatibilityKey: kCFBooleanTrue, kCVPixelBufferCGBitmapContextCompatibilityKey: kCFBooleanTrue] as CFDictionary
      var pixelBuffer : CVPixelBuffer?
      let status = CVPixelBufferCreate(kCFAllocatorDefault, Int(image.size.width), Int(image.size.height), kCVPixelFormatType_32ARGB, attrs, &pixelBuffer)
      guard (status == kCVReturnSuccess) else {
        return nil
      }

      CVPixelBufferLockBaseAddress(pixelBuffer!, CVPixelBufferLockFlags(rawValue: 0))
      let pixelData = CVPixelBufferGetBaseAddress(pixelBuffer!)

      let rgbColorSpace = CGColorSpaceCreateDeviceRGB()
      let context = CGContext(data: pixelData, width: Int(image.size.width), height: Int(image.size.height), bitsPerComponent: 8, bytesPerRow: CVPixelBufferGetBytesPerRow(pixelBuffer!), space: rgbColorSpace, bitmapInfo: CGImageAlphaInfo.noneSkipFirst.rawValue)

      context?.translateBy(x: 0, y: image.size.height)
      context?.scaleBy(x: 1.0, y: -1.0)

      UIGraphicsPushContext(context!)
      image.draw(in: CGRect(x: 0, y: 0, width: image.size.width, height: image.size.height))
      UIGraphicsPopContext()
      CVPixelBufferUnlockBaseAddress(pixelBuffer!, CVPixelBufferLockFlags(rawValue: 0))

      return pixelBuffer
    }
    
    
}

// MARK: - Delegate
extension ViewController:UIImagePickerControllerDelegate {
    
    @available(iOS 2.0, *)
    func imagePickerController(_ picker: UIImagePickerController, didFinishPickingMediaWithInfo info: [UIImagePickerController.InfoKey : Any]) {
        
        guard let image = info[.originalImage] as? UIImage  else {
            fatalError("the image is not choiced.")
        }
           
        self.imageView.image = image

        guard let friutRecog = try? model.prediction(image: buffer(from: image)!) else {
            print("model can not recoginize the image")
            return
        }
        
        nameLabel.text = "Fruit name: \(friutRecog.classLabel)"
        probLabel.text = "Probability: \(friutRecog.classLabelProbs)"
        
        imageController.dismiss(animated: true)
         
    }

    @available(iOS 2.0, *)
    func imagePickerControllerDidCancel(_ picker: UIImagePickerController) {
        print("user is cancel choice image")
        imageController.dismiss(animated: true)
    }
}



